"""
被控模式实现
"""

import platform
import signal
import sys
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from relay_api import RelayAPI
from ssh_manager import SSHManager
from config import RELAY_SERVER_HOST

console = Console()

class AgentMode:
    def __init__(self):
        self.api = RelayAPI()
        self.ssh_manager = SSHManager()
        self.current_key = None
        self.is_running = False
    
    def start(self):
        """启动被控模式"""
        console.print("\n[bold green]启动被控模式...[/bold green]")
        
        # 获取系统信息
        client_info = f"{platform.system()} {platform.release()} - {platform.machine()}"
        
        # 注册到中继服务器
        console.print("正在向中继服务器注册...")
        result = self.api.register(client_info=client_info)
        
        if result.get("status") != "success":
            console.print(f"[bold red]注册失败: {result.get('message', '未知错误')}[/bold red]")
            return False
        
        self.current_key = result["key"]
        tunnel_port = result["tunnel_port"]
        expires_in = result.get("expires_in", 86400)
        
        # 显示连接密钥
        self._display_key_info(self.current_key, expires_in)
        
        # 建立SSH反向隧道
        console.print("正在建立SSH反向隧道...")
        success = self.ssh_manager.create_reverse_tunnel(
            relay_host=RELAY_SERVER_HOST,
            relay_port=22,  # 假设中继服务器SSH端口是22
            tunnel_port=tunnel_port
        )
        
        if not success:
            console.print("[bold red]建立SSH隧道失败[/bold red]")
            return False
        
        # 启动心跳
        self.ssh_manager.start_heartbeat(lambda: self.api.heartbeat(self.current_key))
        
        # 设置信号处理
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)
        
        self.is_running = True
        console.print("[bold green]被控模式已启动，等待连接中...[/bold green]")
        console.print("按 [bold yellow]Ctrl+C[/bold yellow] 停止服务")
        
        # 保持运行状态
        try:
            while self.is_running:
                if not self.ssh_manager.is_tunnel_active():
                    console.print("[bold red]SSH隧道已断开[/bold red]")
                    break
                
                import time
                time.sleep(1)
        except KeyboardInterrupt:
            pass
        
        self.stop()
        return True
    
    def stop(self):
        """停止被控模式"""
        if not self.is_running:
            return
        
        console.print("\n[bold yellow]正在停止被控模式...[/bold yellow]")
        self.is_running = False
        
        # 停止SSH隧道
        self.ssh_manager.stop_tunnel()
        
        # 注销连接
        if self.current_key:
            result = self.api.unregister(self.current_key)
            if result.get("status") == "success":
                console.print("[bold green]已成功注销连接[/bold green]")
            else:
                console.print(f"[bold yellow]注销连接失败: {result.get('message', '未知错误')}[/bold yellow]")
        
        console.print("[bold green]被控模式已停止[/bold green]")
    
    def _display_key_info(self, key: str, expires_in: int):
        """显示密钥信息"""
        hours = expires_in // 3600
        
        key_text = Text(key, style="bold cyan")
        info_text = Text(f"有效期: {hours}小时", style="dim")
        
        panel_content = f"""
[bold white]您的连接密钥是:[/bold white]

[bold cyan]{key}[/bold cyan]

[dim]请将此密钥提供给需要连接的用户
有效期: {hours}小时[/dim]
        """.strip()
        
        panel = Panel(
            panel_content,
            title="[bold green]连接密钥[/bold green]",
            border_style="green",
            padding=(1, 2)
        )
        
        console.print(panel)
    
    def _signal_handler(self, signum, frame):
        """信号处理器"""
        console.print(f"\n[bold yellow]收到信号 {signum}，正在停止...[/bold yellow]")
        self.stop()
        sys.exit(0)

